#include "app_encryption.h"

#include "aes.h"
#include "base64.h"
#include "cJSON.h"
#include "sha256.h"
#include "base64.h"
#include "stdlib.h"
#include "w25q16.h"
#include "http.h"

#define BUFFER_SIZE 256
#define AES_BLOCK_SIZE 16

#define KEY_FLAG 0x000000
#define KEY_ADDR 0x000000 + 1
#define IV_ADDR 0x000000 + 1 + 32

#define KEY_FLAG_LEN 1
#define KEY_LEN 32
#define IV_LEN 16

// 旧密钥
static uint8_t key[KEY_LEN] = {
    0x99, 0xfb, 0xf7, 0xd3, 0x02, 0x7a, 0x39, 0x99, 
    0x4c, 0x61, 0xe7, 0x29, 0x57, 0x44, 0x9c, 0x39,
    0x67, 0xcb, 0xbd, 0x6b, 0x3b, 0xa3, 0x35, 0xf2,
    0x69, 0x44, 0x5b, 0x4f, 0x40, 0x78, 0xb3, 0x2f
};

static uint8_t iv[IV_LEN] = "0000000000000000";

// 用来存放新的密钥暂时固定
static uint8_t new_key[KEY_LEN] = {0};
  
static uint8_t new_iv[IV_LEN] = "0000000000000000";

/**
 * @brief STATE_KEY_ROTATE 状态处理 - 共享密钥更换密钥
 * @return 成功0，失败-1
 */
void Init_Key_AND_Iv()
{
    // uint8_t keyflag = 0; 
    // W25Q16_ReadData(KEY_FLAG, &keyflag, KEY_FLAG_LEN);
    // if(keyflag == 1)
    // {
    //     W25Q16_ReadData(KEY_ADDR, key, KEY_LEN);
    //     // printf("read key: ");
    //     // for (size_t i = 0; i < KEY_LEN; i++) {
    //     //     printf("%02x", key[i]);
    //     // }
    //     // printf("\n");
        
        // W25Q16_ReadData(IV_ADDR, iv, IV_LEN);
    //     // printf("read iv: ");
    //     // for (size_t i = 0; i < IV_LEN; i++) {
    //     //     printf("%02x", key[i]);
    //     // }
    //     // printf("\n");
    // }
}

/**
 * @brief STATE_KEY_ROTATE 状态处理 - 共享密钥更换密钥
 */
uint8_t Handle_Key_Rotate_State(void) {
    char message[] = "3154315432534";
    char keys[] = "432535432";
    uint8_t ret = -1;
    uint8_t keyflag = 1;

    // 用随机函数生成随机值
    uint32_t randvale =  rand();

    keys[0] =  (randvale << 24) & 0xFF; 
    keys[1] = (randvale << 16) & 0xFF; 
    keys[2] = (randvale << 8)  & 0xFF;
    keys[3] = (randvale)  & 0xFF;

    // HMAC-SHA256测试 (更新 key)
    if ((ret = compute_hash(new_key, message, sizeof(message) - 1,
                    keys, sizeof(keys) - 1, HASH_HMAC_SHA256)) == 0)
    {
        // printf("HMAC-SHA256: ");
        // for (size_t i = 0; i < KEY_LEN; i++) {
        //     printf("%02x", new_key[i]);
        // }
        // printf("\n");
    }

    // 更新IV
    // memcpy(new_iv, new_key, IV_LEN);

    /* 写到flash */
    W25Q16_EraseSector(KEY_FLAG); 
  
    // 写入数据到地址
    SPI_FLASH_BufferWrite(&keyflag, KEY_FLAG, KEY_FLAG_LEN);
    SPI_FLASH_BufferWrite(new_key, KEY_ADDR, KEY_LEN);
    SPI_FLASH_BufferWrite(new_iv, IV_ADDR, IV_LEN);
    
    char *base64_key = NULL;
    char *base64_iv = NULL;

    /* 做包 */
    cJSON *root = cJSON_CreateObject();
    if (!root) {
        fprintf(stderr, "Error: Failed to create root object\n");
        return -1;
    }
    
    cJSON_AddStringToObject(root, "newSecretKey", message);
    cJSON_AddStringToObject(root, "newHmacKey", keys);
    cJSON_AddStringToObject(root, "newIV", "0000000000000000");
    
    char *json_str = cJSON_PrintUnformatted(root);
    if (!json_str) {
        fprintf(stderr, "Error: Serialization failed\n");
        cJSON_Delete(root);
        return -1;
    }
    // printf("Successfully serialized JSON:\n%s\n", json_str);
    // printf("json_str = 0%p\n", json_str);
    
    size_t orig_len = strlen(json_str);  
    
    char* result;
    
    encrypt_to_base64(json_str, orig_len, AES_CBC_MODE,  &result);
    
    // printf("Encoded:%s\n", result);
    
    // 直接将新生成的key和Iv推送并加密数据给服务器
    
    cJSON *pload = cJSON_CreateObject();
    if (!pload) {
        fprintf(stderr, "Error: Failed to create root object\n");
        return -1;
    }
    cJSON_AddStringToObject(pload, "encryptedData", result);
    cJSON_AddStringToObject(pload, "encryptionMode", "cbc");

    char *pload_str = cJSON_PrintUnformatted(pload);
    if (!pload_str) {
        fprintf(stderr, "Error: Serialization failed\n");
        cJSON_Delete(pload);
        return -1;
    }

    // printf("The most total:%s\n", pload_str);
    
    // 发送POST请求
    HttpRequest request = {
        .type = HTTP_POST,
        .host = IP_ADDR,
        .path = ADDR_PATH_KEY,
        .json_body = pload_str,
        .port = PORT
    };


    http_request_send(&request);


    // 更新加密方法
    memcpy(key, new_key, KEY_LEN);

    free(json_str);
    cJSON_Delete(root);
    free(result);
    cJSON_Delete(pload);
    free(pload_str);

    return 0; // 未更新时返回无效值
}

/**
 * @brief PKCS#7 添加填充
 * @param data 待填充的数据缓冲区（必须有足够空间容纳填充字节）
 * @param data_len 原始数据长度（不需要是块大小的整数倍）
 * @param block_size 加密块大小（如AES为16）
 * @return 填充后的总长度（等于 data_len + pad_value）
 */
size_t pkcs7_pad(uint8_t *data, size_t data_len, size_t block_size) {
    uint8_t pad_value = block_size - (data_len % block_size);
    memset(data + data_len, pad_value, pad_value);
    return data_len + pad_value;
}
  
/**
 * @brief PKCS#7 去除填充
 * @param data 带填充的数据（必须是块大小的整数倍）
 * @param data_len 数据总长度（必须为块大小的整数倍）
 * @param block_size 加密块大小（如AES为16）
 * @return 去除填充后的实际数据长度（失败返回0）
 */
size_t pkcs7_unpad(const uint8_t *data, size_t data_len, size_t block_size) {
    // 检查输入有效性
    if (data_len == 0 || data_len % block_size != 0) return 0;

    uint8_t pad_value = data[data_len - 1];
    // 验证填充值范围（1到block_size）
    if (pad_value == 0 || pad_value > block_size) return 0;

    // 检查所有填充字节
    for (size_t i = data_len - pad_value; i < data_len; i++) {
        if (data[i] != pad_value) return 0;
    }

    return data_len - pad_value;
}
  
// 通用哈希计算接口
/**
 * @brief 计算哈希值
 * @param output 输出缓冲区（至少32字节）
 * @param input 输入数据
 * @param input_len 输入长度
 * @param key HMAC密钥（仅HMAC模式需要）
 * @param key_len HMAC密钥长度
 * @param algo 算法类型
 * @return 0成功，-1失败
 */
int compute_hash(uint8_t *output, 
                const uint8_t *input, size_t input_len,
                const uint8_t *key, size_t key_len,
                hash_algorithm_t algo)
{
    if (!output || !input) return -1;

    switch (algo) {
        case HASH_SHA256:
            if (key || key_len > 0) {
                printf("[WARN] Key ignored for SHA256\n");
            }
            sha256_get(output, input, input_len);
            break;

        case HASH_HMAC_SHA256:
            if (!key || key_len == 0) {
                printf("[ERROR] HMAC requires key\n");
                return -1;
            }
            hmac_sha256_get(output, (uint8_t *)input, input_len, (uint8_t *)key, key_len);
            break;

        default:
            printf("[ERROR] Unknown algorithm\n");
            return -1;
    }
    return 0;
}

// 示例用法
// void test_hash_interface()
// {
//     uint8_t message[] = "your-secret-key-32";
//     uint8_t key[] = "121354";
    
//     uint8_t digest[32];
    
//     // SHA256测试
//     if (compute_hash(digest, message, sizeof(message)-1, 
//                     NULL, 0, HASH_SHA256) == 0) 
//     {
//         printf("SHA256: ");
//         for (size_t i = 0; i < 32; i++) {
//             printf("%02x", digest[i]);
//         }
//         printf("\n");
//     }

//     // HMAC-SHA256测试
//     if (compute_hash(digest, message, sizeof(message)-1,
//                     key, sizeof(key)-1, HASH_HMAC_SHA256) == 0)
//     {
//         printf("HMAC-SHA256: ");
//         for (size_t i = 0; i < 32; i++) {
//             printf("%02x", digest[i]);
//         }
//         printf("\n");
//     }
// }
/**
 * @brief 统一的AES加密接口
 * @param input 输入数据（必须是16字节的整数倍）
 * @param output 输出缓冲区（大小至少等于输入长度）
 * @param length 数据长度（字节数）
 * @param mode 加密模式（AES_ECB_MODE / AES_CBC_MODE）
 * @return 0成功，-1失败
 */
int aes_encrypt(const uint8_t *input, uint8_t *output, 
    size_t length, aes_mode_t mode)
{
    // 参数检查
    if (!input || !output || length == 0 || length % 16 != 0) {
    return -1;
    }

    struct AES_ctx ctx;
    AES_init_ctx(&ctx, key); // 初始化密钥

    switch (mode) {
        case AES_ECB_MODE:
            // ECB模式加密
            for (size_t i = 0; i < length; i += 16) {
                memcpy(output + i, input + i, 16);
                AES_ECB_encrypt(&ctx, output + i);
            }
        break;

        case AES_CBC_MODE:
            // CBC模式需要设置IV
            memcpy(ctx.Iv, iv, 16);
            memcpy(output, input, length);
            AES_CBC_encrypt_buffer(&ctx, output, length);
        break;

        default:
        return -1;
    }

    return 0;
}

// // 示例用法
// void example_usage() {
//     uint8_t plaintext[32] = { /* 32字节明文数据 */ };
//     uint8_t ciphertext[32];

//     // ECB模式加密
//     if (aes_encrypt(plaintext, ciphertext, sizeof(plaintext), AES_ECB_MODE) == 0) {
//         // 加密成功处理
//     }

//     // CBC模式加密
//     if (aes_encrypt(plaintext, ciphertext, sizeof(plaintext), AES_CBC_MODE) == 0) {
//         // 加密成功处理
//     }
// }

/**
 * @brief 统一的AES解密接口
 * @param input 输入密文（必须是16字节的整数倍）
 * @param output 输出缓冲区（大小至少等于输入长度）
 * @param length 数据长度（字节数）
 * @param mode 解密模式（AES_ECB_MODE / AES_CBC_MODE）
 * @return 0成功，-1失败
 */
int aes_decrypt(const uint8_t *input, uint8_t *output, 
    size_t length, aes_mode_t mode)
{
    // 参数安全检查
    if (!input || !output || length == 0 || length % 16 != 0) {
        return -1; // AES要求16字节对齐
    }

    struct AES_ctx ctx;
    AES_init_ctx(&ctx, key); // 初始化密钥

    switch (mode) {
        case AES_ECB_MODE:
            // ECB模式解密
            for (size_t i = 0; i < length; i += 16) {
                memcpy(output + i, input + i, 16);
                AES_ECB_decrypt(&ctx, output + i);
            }
        break;

        case AES_CBC_MODE:
            // CBC模式需要设置IV
            memcpy(ctx.Iv, iv, 16);
            memcpy(output, input, length);
            AES_CBC_decrypt_buffer(&ctx, output, length);
        break;

        default:
        return -1;
    }

    return 0;
}

// // 示例用法
// void example_usage() {
//     uint8_t ciphertext[32] = { /* 32字节密文 */ };
//     uint8_t plaintext[32];

//     // ECB模式解密
//     if (aes_decrypt(ciphertext, plaintext, sizeof(ciphertext), AES_ECB_MODE) == 0) {
//         // 解密成功处理
//     }

//     // CBC模式解密（带错误码检查）
//     aes_result_t ret = aes_decrypt_ex(ciphertext, plaintext, 
//                                     sizeof(ciphertext), AES_CBC_MODE);
//     if (ret == AES_SUCCESS) {
//         // 解密成功处理
//     } else {
//         // 处理错误
//         printf("解密失败，错误码: %d\n", ret);
//     }
// }

/**
 * @brief 完整的加密->Base64流程（支持CBC/ECB模式）
 * @param plaintext 原始明文数据
 * @param plaintext_len 明文长度
 * @param mode 加密模式（AES_CBC_MODE / AES_ECB_MODE）
 * @param base64_output 输出Base64字符串指针（需要外部释放）
 * @return 0成功，-1失败
 */
int encrypt_to_base64(const uint8_t *plaintext, size_t plaintext_len, int mode, char **base64_output) 
{
    if (!plaintext || !base64_output || plaintext_len == 0) {
        return -1;
    }

    uint8_t buffer[BUFFER_SIZE] = {0};
    uint8_t encrypted[BUFFER_SIZE] = {0};
    size_t padded_len;
    size_t b64_len;

    // 1. 拷贝并填充数据
    if (plaintext_len > BUFFER_SIZE - AES_BLOCK_SIZE) {
        printf("输入数据过长\n");
        return -1;
    }
    memcpy(buffer, plaintext, plaintext_len);
    padded_len = pkcs7_pad(buffer, plaintext_len, AES_BLOCK_SIZE);

    // 2. AES加密（动态选择模式）
    if (aes_encrypt(buffer, encrypted, padded_len, mode) != 0) {
        printf("加密失败\n");
        return -1;
    }

    // 3. Base64编码
    *base64_output = base64_encode(encrypted, padded_len, &b64_len);
    if (!*base64_output) {
        printf("Base64编码失败\n");
        return -1;
    }

    return 0;
}

/**
 * @brief 完整的Base64->解密流程（支持CBC/ECB模式）
 * @param base64_input Base64编码字符串
 * @param mode 解密模式（AES_CBC_MODE / AES_ECB_MODE）
 * @param output 输出缓冲区
 * @param output_size 输出缓冲区大小
 * @return 解密后的数据长度，-1表示失败
 */
int base64_to_decrypt(const char *base64_input, int mode, uint8_t *output, size_t output_size) 
{
    if (!base64_input || !output) return -1;

    uint8_t decoded[BUFFER_SIZE] = {0};
    uint8_t decrypted[BUFFER_SIZE] = {0};
    size_t decoded_len;
    int plaintext_len;

    // 1. Base64解码
    unsigned char *tmp = base64_decode(base64_input, strlen(base64_input), &decoded_len);
    if (!tmp) {
        printf("Base64解码失败\n");
        return -1;
    }
    memcpy(decoded, tmp, decoded_len);
    free(tmp);

    // 2. AES解密（动态选择模式）
    if (aes_decrypt(decoded, decrypted, decoded_len, mode) != 0) {
        printf("解密失败\n");
        return -1;
    }

    // 3. 去除填充（仅CBC模式需要，ECB模式通常也需要填充）
    plaintext_len = pkcs7_unpad(decrypted, decoded_len, AES_BLOCK_SIZE);
    printf("plaintext_len%d\n", plaintext_len);
    if (plaintext_len <= 0 || plaintext_len > output_size) {
        printf("输出缓冲区不足或填充错误\n");
        return -1;
    }

    memcpy(output, decrypted, plaintext_len);
    return plaintext_len;
}